# middleware/redis_client.py
from typing import Any, Optional, Dict, Union
from redis import Redis, ConnectionPool
from redis.lock import Lock  # Redis官方Lock，无第三方依赖
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
from loguru import logger
import json
import hashlib

# 导入自定义工具和常量
from utils.logger import Logger
from utils.exceptions import RedisError, LockAcquireError
from utils.constants import (
    REDIS_HOST,
    REDIS_PORT,
    REDIS_DB,
    REDIS_PASSWORD,
    REDIS_POOL_SIZE,
    REDIS_LOCK_TIMEOUT,
    REDIS_LOCK_RETRIES,
    REDIS_LOCK_RETRY_DELAY,
    REDIS_RETRY_ATTEMPTS,
    REDIS_RETRY_WAIT_SECONDS
)

# 初始化日志（固定agent_name为redis_client）
logger = Logger.get_logger(agent_name="redis_client")

class RedisClient:
    """Redis客户端封装（单例模式，兼容tenacity新版本）"""
    _instance = None
    _lock_instances: Dict[str, Lock] = {}  # 存储已获取的锁实例

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(RedisClient, cls).__new__(cls)
            cls._instance._init_connection()
        return cls._instance

    def _init_connection(self):
        """初始化Redis连接池"""
        try:
            pool = ConnectionPool(
                host=REDIS_HOST,
                port=REDIS_PORT,
                db=REDIS_DB,
                password=REDIS_PASSWORD,
                max_connections=REDIS_POOL_SIZE,
                decode_responses=False  # 支持二进制数据，手动解码
            )
            self.client = Redis(connection_pool=pool)
            self.client.ping()  # 测试连接
            logger.success(f"Redis连接池初始化成功，DB={REDIS_DB}，连接池大小={REDIS_POOL_SIZE}")
        except Exception as e:
            logger.error(f"Redis连接初始化失败: {str(e)}")
            raise RuntimeError("Redis连接初始化失败，服务无法启动") from e

    # ---------------------- 基础String操作（修复tenacity兼容） ----------------------
    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args（tenacity新版本属性名）
        before=lambda retry_state: logger.warning(f"Redis GET重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}"),
        reraise=True
    )
    def get(self, key: str) -> Optional[Union[str, bytes]]:
        """获取缓存值（支持str/bytes，自动解码str）"""
        try:
            value = self.client.get(key)
            if isinstance(value, bytes):
                return value.decode("utf-8")
            return value
        except Exception as e:
            logger.error(f"Redis GET失败 key={key}: {str(e)}")
            raise RedisError(message="Redis GET failed", context={"key": key, "error": str(e)}) from e

    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args
        before=lambda retry_state: logger.warning(f"Redis SET重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}"),
        reraise=True
    )
    def set(self, key: str, value: Union[str, bytes], ex: Optional[int] = None) -> bool:
        """设置缓存（支持str/bytes，自动编码）"""
        try:
            if isinstance(value, str):
                value = value.encode("utf-8")
            result = self.client.set(key, value, ex=ex)
            return bool(result)
        except Exception as e:
            logger.error(f"Redis SET失败 key={key}: {str(e)}")
            raise RedisError(message="Redis SET failed", context={"key": key, "error": str(e)}) from e

    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args
        before=lambda retry_state: logger.warning(f"Redis DELETE重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}"),
        reraise=True
    )
    def delete(self, key: str) -> int:
        """删除缓存（返回删除的key数量）"""
        try:
            return self.client.delete(key)
        except Exception as e:
            logger.error(f"Redis DELETE失败 key={key}: {str(e)}")
            raise RedisError(message="Redis DELETE failed", context={"key": key, "error": str(e)}) from e

    # ---------------------- 关键：Hash操作（修复tenacity兼容） ----------------------
    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args，key是args[1]，field是args[2]
        before=lambda retry_state: logger.warning(f"Redis HGET重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}, field={retry_state.args[2]}"),
        reraise=True
    )
    def hget(self, key: str, field: str) -> Optional[Union[str, bytes]]:
        """获取Hash字段值"""
        try:
            value = self.client.hget(key, field)
            if isinstance(value, bytes):
                return value.decode("utf-8")
            return value
        except Exception as e:
            logger.error(f"Redis HGET失败 key={key}, field={field}: {str(e)}")
            raise RedisError(message="Redis HGET failed", context={"key": key, "field": field, "error": str(e)}) from e

    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args
        before=lambda retry_state: logger.warning(f"Redis HSET重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}, field={retry_state.args[2]}"),
        reraise=True
    )
    def hset(self, key: str, field: str, value: Union[str, bytes]) -> bool:
        """设置Hash字段值"""
        try:
            if isinstance(value, str):
                value = value.encode("utf-8")
            result = self.client.hset(key, field, value)
            return bool(result)
        except Exception as e:
            logger.error(f"Redis HSET失败 key={key}, field={field}: {str(e)}")
            raise RedisError(message="Redis HSET failed", context={"key": key, "field": field, "error": str(e)}) from e

    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args
        before=lambda retry_state: logger.warning(f"Redis HGETALL重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}"),
        reraise=True
    )
    def hgetall(self, key: str) -> Dict[str, str]:
        """获取Hash所有字段和值（返回str字典，自动解码）"""
        try:
            data = self.client.hgetall(key)
            return {k.decode("utf-8"): v.decode("utf-8") for k, v in data.items()} if data else {}
        except Exception as e:
            logger.error(f"Redis HGETALL失败 key={key}: {str(e)}")
            raise RedisError(message="Redis HGETALL failed", context={"key": key, "error": str(e)}) from e

    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args
        before=lambda retry_state: logger.warning(f"Redis HSETALL重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}"),
        reraise=True
    )
    def hsetall(self, key: str, mapping: Dict[str, Union[str, bytes]]) -> bool:
        """批量设置Hash字段值"""
        try:
            encoded_mapping = {}
            for k, v in mapping.items():
                encoded_k = k.encode("utf-8") if isinstance(k, str) else k
                encoded_v = v.encode("utf-8") if isinstance(v, str) else v
                encoded_mapping[encoded_k] = encoded_v
            self.client.hset(key, mapping=encoded_mapping)
            return True
        except Exception as e:
            logger.error(f"Redis HSETALL失败 key={key}: {str(e)}")
            raise RedisError(message="Redis HSETALL failed", context={"key": key, "error": str(e)}) from e

    # ---------------------- 分布式锁（Redis官方实现） ----------------------
    def acquire_lock(self, lock_name: str, timeout: int = REDIS_LOCK_TIMEOUT) -> bool:
        """获取分布式锁（支持重试和超时释放）"""
        try:
            if lock_name in self._lock_instances:
                logger.warning(f"锁{lock_name}已被当前实例持有，无需重复获取")
                return True

            lock = Lock(
                redis=self.client,
                name=lock_name,
                timeout=timeout,
                blocking=True,
                blocking_timeout=REDIS_LOCK_RETRIES * REDIS_LOCK_RETRY_DELAY
            )

            acquire_success = False
            for attempt in range(REDIS_LOCK_RETRIES):
                if lock.acquire():
                    acquire_success = True
                    break
                logger.warning(f"获取锁{lock_name}失败，第{attempt+1}次重试（共{REDIS_LOCK_RETRIES}次）")
                import time
                time.sleep(REDIS_LOCK_RETRY_DELAY)

            if acquire_success:
                self._lock_instances[lock_name] = lock
                logger.debug(f"获取分布式锁成功 lock={lock_name}, timeout={timeout}s")
                return True
            logger.error(f"获取分布式锁失败（已重试{REDIS_LOCK_RETRIES}次） lock={lock_name}")
            raise LockAcquireError(message="Redis lock acquire failed", context={"lock_name": lock_name, "retries": REDIS_LOCK_RETRIES})
        except LockAcquireError:
            raise
        except Exception as e:
            logger.error(f"Redis锁获取失败 lock={lock_name}: {str(e)}")
            raise RedisError(message="Redis lock acquire failed", context={"lock_name": lock_name, "error": str(e)}) from e

    def release_lock(self, lock_name: str) -> bool:
        """释放分布式锁"""
        try:
            lock = self._lock_instances.get(lock_name)
            if not lock:
                logger.warning(f"无已持有锁{lock_name}，无需释放")
                return True

            if lock.locked():
                lock.release()
            del self._lock_instances[lock_name]
            logger.debug(f"释放分布式锁成功 lock={lock_name}")
            return True
        except Exception as e:
            logger.error(f"Redis锁释放失败 lock={lock_name}: {str(e)}")
            raise RedisError(message="Redis lock release failed", context={"lock_name": lock_name, "error": str(e)}) from e

    def release_all_locks(self) -> None:
        """释放所有持有的锁"""
        for lock_name in list(self._lock_instances.keys()):
            self.release_lock(lock_name)
        logger.info("已释放所有分布式锁")

    # ---------------------- Pub/Sub（支持优雅退出） ----------------------
    def publish(self, channel: str, message: Union[str, Dict[str, Any]]) -> int:
        """发布消息（支持字符串/字典）"""
        try:
            if isinstance(message, dict):
                message = json.dumps(message, ensure_ascii=False)
            if isinstance(message, str):
                message = message.encode("utf-8")
            subscriber_count = self.client.publish(channel, message)
            logger.debug(f"发布消息成功 channel={channel}, subscriber_count={subscriber_count}")
            return subscriber_count
        except Exception as e:
            logger.error(f"Redis发布失败 channel={channel}: {str(e)}")
            raise RedisError(message="Redis publish failed", context={"channel": channel, "error": str(e)}) from e

    def subscribe(self, channel: str, callback: callable, stop_event) -> None:
        """订阅频道（支持优雅退出）"""
        pubsub = self.client.pubsub()
        pubsub.subscribe(channel)
        logger.info(f"开始订阅频道: {channel}")

        try:
            while not stop_event.is_set():
                message = pubsub.get_message(timeout=1)
                if message and message['type'] == 'message':
                    try:
                        data = message['data']
                        if isinstance(data, bytes):
                            data = data.decode("utf-8")
                        try:
                            data = json.loads(data)
                        except json.JSONDecodeError:
                            pass
                        callback(data)
                    except Exception as e:
                        logger.error(f"消息处理失败 channel={channel}, data={data}: {str(e)}")
        except Exception as e:
            logger.error(f"订阅频道失败 channel={channel}: {str(e)}")
            raise RedisError(message="Redis subscribe failed", context={"channel": channel, "error": str(e)}) from e
        finally:
            pubsub.unsubscribe(channel)
            logger.info(f"已取消订阅频道: {channel}")

    # ---------------------- 过期时间操作 ----------------------
    @retry(
        stop=stop_after_attempt(REDIS_RETRY_ATTEMPTS),
        wait=wait_exponential(multiplier=1, min=REDIS_RETRY_WAIT_SECONDS, max=10),
        retry=retry_if_exception_type((ConnectionRefusedError, TimeoutError)),
        # 关键修复：fn_args → args
        before=lambda retry_state: logger.warning(f"Redis EXPIRE重试中，第{retry_state.attempt_number}次 key={retry_state.args[1]}"),
        reraise=True
    )
    def expire(self, key: str, seconds: int) -> bool:
        """设置key的过期时间"""
        try:
            return self.client.expire(key, seconds)
        except Exception as e:
            logger.error(f"Redis EXPIRE失败 key={key}, seconds={seconds}: {str(e)}")
            raise RedisError(message="Redis expire failed", context={"key": key, "seconds": seconds, "error": str(e)}) from e

    def keys(self, pattern: str) -> list:
        """根据模式匹配获取key列表"""
        try:
            keys = self.client.keys(pattern)
            return [key.decode("utf-8") if isinstance(key, bytes) else key for key in keys]
        except Exception as e:
            logger.error(f"Redis KEYS失败 pattern={pattern}: {str(e)}")
            raise RedisError(message="Redis keys failed", context={"pattern": pattern, "error": str(e)}) from e

# 使用示例（测试所有核心功能）
if __name__ == "__main__":
    import sys
    import time
    import threading
    sys.path.append("..")  # 添加项目根目录到Python路径

    # 临时定义常量（如果utils/constants.py未生效，取消注释）
    # class TempConstants:
    #     REDIS_HOST = "localhost"
    #     REDIS_PORT = 6379
    #     REDIS_DB = 0
    #     REDIS_PASSWORD = ""
    #     REDIS_POOL_SIZE = 10
    #     REDIS_LOCK_TIMEOUT = 5
    #     REDIS_LOCK_RETRIES = 3
    #     REDIS_LOCK_RETRY_DELAY = 0.1
    #     REDIS_RETRY_ATTEMPTS = 3
    #     REDIS_RETRY_WAIT_SECONDS = 1
    #
    # # 替换全局常量
    # global REDIS_HOST, REDIS_PORT, REDIS_DB, REDIS_PASSWORD, REDIS_POOL_SIZE
    # global REDIS_LOCK_TIMEOUT, REDIS_LOCK_RETRIES, REDIS_LOCK_RETRY_DELAY
    # global REDIS_RETRY_ATTEMPTS, REDIS_RETRY_WAIT_SECONDS
    #
    # REDIS_HOST = TempConstants.REDIS_HOST
    # REDIS_PORT = TempConstants.REDIS_PORT
    # REDIS_DB = TempConstants.REDIS_DB
    # REDIS_PASSWORD = TempConstants.REDIS_PASSWORD
    # REDIS_POOL_SIZE = TempConstants.REDIS_POOL_SIZE
    # REDIS_LOCK_TIMEOUT = TempConstants.REDIS_LOCK_TIMEOUT
    # REDIS_LOCK_RETRIES = TempConstants.REDIS_LOCK_RETRIES
    # REDIS_LOCK_RETRY_DELAY = TempConstants.REDIS_LOCK_RETRY_DELAY
    # REDIS_RETRY_ATTEMPTS = TempConstants.REDIS_RETRY_ATTEMPTS
    # REDIS_RETRY_WAIT_SECONDS = TempConstants.REDIS_RETRY_WAIT_SECONDS

    # 初始化RedisClient
    redis_client = RedisClient()

    # 1. 测试String操作
    redis_client.set("test_string_key", "测试值", ex=60)
    print("String GET:", redis_client.get("test_string_key"))  # 输出：测试值

    # 2. 测试Hash操作
    hash_key = "test_hash_key"
    redis_client.hsetall(hash_key, {
        "field1": "value1",
        "field2": json.dumps({"key": "value"}, ensure_ascii=False)
    })
    print("Hash HGETALL:", redis_client.hgetall(hash_key))  # 输出：{'field1': 'value1', 'field2': '{"key":"value"}'}
    print("Hash HGET field2:", redis_client.hget(hash_key, "field2"))  # 输出：{"key":"value"}

    # 3. 测试分布式锁
    print("获取锁:", redis_client.acquire_lock("test_lock"))  # 输出：True
    print("重复获取锁:", redis_client.acquire_lock("test_lock"))  # 输出：True
    print("释放锁:", redis_client.release_lock("test_lock"))  # 输出：True
    print("释放不存在的锁:", redis_client.release_lock("test_lock"))  # 输出：True

    # 4. 测试发布订阅
    stop_event = threading.Event()

    def test_callback(data):
        print("收到订阅消息:", data)

    subscribe_thread = threading.Thread(
        target=redis_client.subscribe,
        args=("test_channel", test_callback, stop_event),
        daemon=True
    )
    subscribe_thread.start()

    redis_client.publish("test_channel", "普通字符串消息")
    redis_client.publish("test_channel", {"type": "json消息", "content": "测试"})
    time.sleep(2)

    stop_event.set()
    subscribe_thread.join()
    print("订阅线程已退出")

    # 5. 释放所有锁
    redis_client.release_all_locks()